from PIL import Image, ImageEnhance
import pytesseract


def enhance_image(image_path):
    """增强图像对比度，有助于提高OCR识别精度"""
    img = Image.open(image_path)
    enhancer = ImageEnhance.Contrast(img)  # 对比度增强
    img_enhanced = enhancer.enhance(2.0)  # 提高对比度因子
    return img_enhanced


def recognize_text_from_image(image_path, tessdata_dir=None):
    """
    从给定的图片路径中提取文字信息。

    :param image_path: 图片文件路径
    :param tessdata_dir: Tesseract的数据目录路径，如果需要的话
    :return: 提取的文字内容
    """
    if tessdata_dir:
        pytesseract.pytesseract.tesseract_cmd = r'[你的Tesseract安装路径]'  # 如果Tesseract不在环境变量中，请指定完整路径

    # 增强图像
    img = enhance_image(image_path)

    # 使用pytesseract进行文字识别
    text = pytesseract.image_to_string(img, lang='chi_sim')  # 'chi_sim'表示简体中文
    return text


def clean_and_join_lines(text):
    """
    清理并拼接提取的文字中的换行符。

    :param text: 提取出的文字内容
    :return: 拼接后更易读的文字内容
    """
    lines = text.split('\n')
    cleaned_lines = []
    current_line = ""

    for line in lines:
        stripped_line = line.strip()
        if stripped_line:
            if current_line:
                current_line += stripped_line
            else:
                current_line = stripped_line
        else:
            if current_line:
                cleaned_lines.append(current_line)
                current_line = ""

    # 处理最后一行
    if current_line:
        cleaned_lines.append(current_line)

    return "\n".join(cleaned_lines)


def extract_address_info(text):
    """
    尝试从提取的文字中分离出收件人的姓名和地址信息。

    :param text: 提取出的文字内容
    :return: 包含收件人和地址的字典
    """
    lines = text.split('\n')
    address_info = {'收件人': '未找到', '地址': '未找到'}

    for i, line in enumerate(lines):
        line = line.strip()
        if "收件人" in line or "收货人" in line:
            address_info['收件人'] = line.replace("收件人:", "").replace("收货人:", "").strip()
        elif "地址" in line:
            address_info['地址'] = line.replace("地址:", "").strip()
        else:
            # 尝试通过上下文分析来猜测收件人和地址
            if not address_info['收件人'] and any(keyword in line for keyword in ['先生', '女士', '小姐']):
                address_info['收件人'] = line.strip()
            if not address_info['地址'] and any(keyword in line for keyword in ['市', '区', '路', '号']):
                address_info['地址'] = line.strip()

    # 如果仍然没有找到，尝试进一步分析
    if address_info['收件人'] == '未找到' or address_info['地址'] == '未找到':
        for i, line in enumerate(lines):
            if address_info['收件人'] == '未找到':
                if any(keyword in line for keyword in ['先生', '女士', '小姐']):
                    address_info['收件人'] = line.strip()
            if address_info['地址'] == '未找到':
                if any(keyword in line for keyword in ['市', '区', '路', '号']):
                    address_info['地址'] = line.strip()

    return address_info


def main():
    # 图片路径，请替换为实际路径
    image_path = '/home/weiqiangren/Python/project/py-basic-exercises/图片文字识别/企业微信截图_1738913734328.png'  # 替换为你的图片路径

    # 调用函数并打印结果
    recognized_text = recognize_text_from_image(image_path)
    print("原始识别文字：")
    print(recognized_text)

    cleaned_text = clean_and_join_lines(recognized_text)
    print("\n清理并拼接后的文字：")
    print(cleaned_text)

    address_info = extract_address_info(cleaned_text)
    print("\n提取的收件人和地址信息：")
    print(address_info)


if __name__ == "__main__":
    main()